import pandas as pd
import numpy as np


def generate_similar_data(csv_path, output_path, method='noise', params=None):
    """
    生成与原始CSV文件相似但有差异的数据

    参数:
    - csv_path: 原始CSV文件路径
    - output_path: 输出CSV文件路径
    - method: 扰动方法，可选 'noise'(噪声), 'scale'(缩放), 'shift'(平移), 'mixed'(混合)
    - params: 方法特定参数
    """
    # 读取原始数据
    df = pd.read_csv(csv_path)

    # 复制一份新数据
    new_df = df.copy()

    # 默认参数
    if params is None:
        params = {}

    # 对数值列应用扰动
    for col in df.select_dtypes(include=np.number).columns:
        data = df[col].values
        std = data.std()

        if method == 'noise':
            # 添加随机噪声
            noise_level = params.get('noise', 0.05)
            noise = np.random.uniform(-noise_level, noise_level, size=len(data)) * data
            new_df[col] = data + noise

        elif method == 'scale':
            # 缩放数据
            scale_range = params.get('scale', (0.95, 1.05))
            scale_factors = np.random.uniform(scale_range[0], scale_range[1], size=len(data))
            new_df[col] = data * scale_factors

        elif method == 'shift':
            # 平移数据
            shift_range = params.get('shift', (-0.05, 0.05))
            shift = np.random.uniform(shift_range[0], shift_range[1], size=len(data)) * std
            new_df[col] = data + shift

        elif method == 'mixed':
            # 混合扰动（随机选择方法）
            method_choice = np.random.choice(['noise', 'scale', 'shift'])
            if method_choice == 'noise':
                noise_level = params.get('noise', 0.05)
                noise = np.random.uniform(-noise_level, noise_level, size=len(data)) * data
                new_df[col] = data + noise
            elif method_choice == 'scale':
                scale_range = params.get('scale', (0.95, 1.05))
                scale_factors = np.random.uniform(scale_range[0], scale_range[1], size=len(data))
                new_df[col] = data * scale_factors
            else:  # shift
                shift_range = params.get('shift', (-0.05, 0.05))
                shift = np.random.uniform(shift_range[0], shift_range[1], size=len(data)) * std
                new_df[col] = data + shift

    # 保存新数据
    new_df.to_csv(output_path, index=False)
    return new_df


if __name__ == '__main__':
    data_list = ['pressure', 'rain', 'rainfall_probability', 'temperature', 'wind_speed']
    area_list = ['中西区', '湾仔区', '东区', '南区', '油尖旺区', '深水埗区', '九龙城区', '黄大仙区', '观塘区', '葵青区',
                 '荃湾区']
    for i in data_list:
        for j in area_list:
            csv_path = f'./data/{i}/北区.csv'
            output_path = f'./data/{i}/{j}_new.csv'
            params = {
                'noise': 0.05,  # 噪声水平
                'scale': (0.95, 1.05),  # 缩放范围
                'shift': (-0.05, 0.05)  # 平移范围
            }
            new_df = generate_similar_data(csv_path, output_path, method='mixed', params=params)
            print(f"生成了 {output_path}，包含与 {csv_path} 相似但有差异的数据。")
